import os
from typing import Dict, Any
from stable_baselines3 import PPO


# 管理单个子策略的策略模型
class PolicyManager:
    def __init__(self, sub_strategy_types: str = None):
        # 将模型存储文件夹中的模型根据子策略全部读取
        self.policy_files_path = os.path.join("./data/trained_models/sub_strategies", sub_strategy_types)

    def upadte_policy(self, sub_strategy_types: str = None):
        # 更新策略文件路径
        self.policy_files_path = os.path.join("./data/trained_models/sub_strategies", sub_strategy_types)

    def get_random_policy(self) -> PPO:
        """        
        随机采样策略
        :return: PPO模型实例
        """
        # 读取全部策略文件
        self.policy_files_path = os.path.join(self.policy_files_path,"models")
        policy_files = [f for f in os.listdir(self.policy_files_path) if f.endswith('.zip')]

        if not policy_files:
            raise ValueError("No policy files found in the specified directory.")
        
        import random
        random_policy_file = random.choice(policy_files)
        policy_path = os.path.join(self.policy_files_path, random_policy_file)

        # 从文件加载策略
        loaded_policy = PPO.load(policy_path)
        return loaded_policy

    def get_latest_policy(self) -> PPO:
        """
        获取最新的策略
        :param plane_id: 飞机ID
        :return: PPO模型实例
        """
        self.policy_files_path = os.path.join(self.policy_files_path,"models")
        policy_files = [f for f in os.listdir(self.policy_files_path) if f.endswith('.zip')]
        if not policy_files:
            raise ValueError("No policy files found in the specified directory.")
        latest_policy_file = policy_files[-1]  # 获取最新的策略文件
        policy_path = os.path.join(self.policy_files_path, latest_policy_file)
        # 从文件加载策略
        loaded_policy = PPO.load(policy_path)
        return loaded_policy

    def get_best_policy(self) -> PPO:
        """
        获取最佳策略
        :param plane_id: 飞机ID
        :return: PPO模型实例
        """
        self.policy_files_path = os.path.join(self.policy_files_path,"best_model")
        policy_files = [f for f in os.listdir(self.policy_files_path) if f.endswith('.zip')]

        if not policy_files:
            raise ValueError("No policy files found in the specified directory.")
        
        latest_policy_file = policy_files[-1]  
        policy_path = os.path.join(self.policy_files_path, latest_policy_file)

        # 从文件加载策略
        loaded_policy = PPO.load(policy_path)
        return loaded_policy
